﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;

namespace EvoBoost.Core
{
	public class ExponentialLossFunction : ILossFunction
	{
		private const double MinusInf = -10;

		private const double PlusInf = 10;

		public double Calculate(double expectedOutcome, double realOutcome)
		{
			double result = Math.Exp(-expectedOutcome * realOutcome);
			Debug.Assert(!Double.IsNaN(result) && !Double.IsInfinity(result));
			return result;
		}

		public double Minimize(IList<double> previouslyBoostedOutcomes, IList<double> expectedOutcomes)
		{
			if (previouslyBoostedOutcomes == null) throw new ArgumentNullException("previouslyBoostedOutcomes");
			if (expectedOutcomes == null) throw new ArgumentNullException("expectedOutcomes");
			if (previouslyBoostedOutcomes.Count != expectedOutcomes.Count)
				throw new ArgumentException(
					"Previously boosted outcome collection should have the same size as expected outcome collection");
			if (previouslyBoostedOutcomes.Count == 0)
				throw new ArgumentException("At least one sample should be provided.");

			double numerator = 0, denominator = 0;
			int sampleCount = previouslyBoostedOutcomes.Count;
			for (int i = 0; i < sampleCount; ++i)
			{
				double weight = Math.Exp(-previouslyBoostedOutcomes[i] * expectedOutcomes[i]);
				if (expectedOutcomes[i] == 1)
					numerator += weight;
				else if (expectedOutcomes[i] == -1)
					denominator += weight;
				else
					throw new InvalidOperationException("Exponential loss should be used only for 2-class classification with -1 and 1 classes.");
			}

			Debug.Assert(numerator >= 0 && denominator >= 0);
			return HalfLogOddsOrInfinity(numerator, denominator);
		}

		private static double HalfLogOddsOrInfinity(double numerator, double denominator)
		{
			if (numerator < Double.Epsilon)
				return MinusInf;
			if (denominator < Double.Epsilon)
				return PlusInf;

			return 0.5 * (Math.Log(numerator) - Math.Log(denominator));
		}

		private static double WeightSumForClass(IEnumerable<SplitterClassPair> splitters, double classValue)
		{
			Debug.Assert(classValue == -1 || classValue == 1);
			
			double weightSum = 0;
			foreach (SplitterClassPair splitter in splitters)
			{
				if (splitter.ClassValue == classValue)
					weightSum += splitter.Weight;
			}

			return weightSum;
		}

		public double FindOptimalSplit(IList<double> previouslyBoostedOutcomes, IList<double> expectedOutcomes, IList<double> splitterValues)
		{
			if (previouslyBoostedOutcomes == null) throw new ArgumentNullException("previouslyBoostedOutcomes");
			if (expectedOutcomes == null) throw new ArgumentNullException("expectedOutcomes");
			if (splitterValues == null) throw new ArgumentNullException("splitterValues");
			if (previouslyBoostedOutcomes.Count != expectedOutcomes.Count || splitterValues.Count != expectedOutcomes.Count)
				throw new ArgumentException("Outcome and splitter value collections should have equal item count.");

			int sampleCount = expectedOutcomes.Count;

			// Sort splitter values together with class info
			List<SplitterClassPair> sortedSplitters = new List<SplitterClassPair>(sampleCount);
			for (int i = 0; i < sampleCount; ++i)
			{
				if (expectedOutcomes[i] != -1 && expectedOutcomes[i] != 1)
					throw new InvalidOperationException("Exponential loss should be used only for 2-class classification with -1 and 1 classes.");
				sortedSplitters.Add(new SplitterClassPair(splitterValues[i], expectedOutcomes[i], previouslyBoostedOutcomes[i]));
			}
			sortedSplitters.Sort((s1, s2) => Math.Sign(s1.Splitter - s2.Splitter));

			double numeratorLeft = WeightSumForClass(sortedSplitters, 1);
			double denominatorLeft = WeightSumForClass(sortedSplitters, -1);
			double classLeft = HalfLogOddsOrInfinity(numeratorLeft, denominatorLeft), classRight = 0;
			double numeratorRight = 0;
			double denominatorRight = 0;
			double loss = sortedSplitters.Sum(s => this.Calculate(s.ClassValue, classLeft + s.AlreadyBoosted));
			double bestLoss = loss;
			double bestThreshold = sortedSplitters[sampleCount - 1].Splitter;
			int index = sampleCount - 1;

			while (index >= 0)
			{
				double numeratorRightOld = numeratorRight, denominatorRightOld = denominatorRight;
				int indexOld = index;

				while (true)
				{
					SplitterClassPair s = sortedSplitters[index];
					if (s.ClassValue == 1)
					{
						numeratorLeft -= s.Weight;
						numeratorRight += s.Weight;
					}
					else if (s.ClassValue == -1)
					{
						denominatorLeft -= s.Weight;
						denominatorRight += s.Weight;
					}
					else
						Debug.Fail("We should never get there.");

					if (index == 0 || sortedSplitters[index].Splitter != sortedSplitters[index - 1].Splitter)
						break;

					index -= 1;
				}

				// Current element is the first from left
				index -= 1;

				double classLeftOld = classLeft, classRightOld = classRight;
				classLeft = HalfLogOddsOrInfinity(numeratorLeft, denominatorLeft);
				classRight = HalfLogOddsOrInfinity(numeratorRight, denominatorRight);

				double deltaLoss = 0;
				deltaLoss += numeratorLeft * (Math.Exp(-classLeftOld) - Math.Exp(-classLeft));
				deltaLoss += denominatorLeft * (Math.Exp(classLeftOld) - Math.Exp(classLeft));
				deltaLoss += numeratorRightOld * (Math.Exp(-classRightOld) - Math.Exp(-classRight));
				deltaLoss += denominatorRightOld * (Math.Exp(classRightOld) - Math.Exp(classRight));
				for (int i = index + 1; i <= indexOld; ++i)
				{
				    SplitterClassPair s = sortedSplitters[i];
				    deltaLoss += s.Weight * (Math.Exp(-s.ClassValue * classLeftOld) - Math.Exp(-s.ClassValue * classRight));
				}
				loss -= deltaLoss;

				if (loss < bestLoss)
				{
					bestLoss = loss;
					bestThreshold = index >= 0 ? sortedSplitters[index].Splitter : sortedSplitters[0].Splitter - 1;
				}
			}

			return bestThreshold;
		}

		public override string ToString()
		{
			return "EXPLOSS";
		}

		private struct SplitterClassPair
		{
			public double Splitter { get; private set; }

			public double ClassValue { get; private set; }

			public double AlreadyBoosted { get; private set; }

			public double Weight { get; private set; }

			public SplitterClassPair(double splitter, double classValue, double alreadyBoosted)
				: this()
			{
				Debug.Assert(classValue == -1 || classValue == 1);

				this.Splitter = splitter;
				this.ClassValue = classValue;
				this.AlreadyBoosted = alreadyBoosted;
				this.Weight = Math.Exp(-this.ClassValue * this.AlreadyBoosted);
			}
		}
	}
}
